/*========================================================================
   Copyright (c) Ars Electronica Futurelab, 2008
  ========================================================================*/

#include "TCPManager.h"

#include <stdio.h>


#ifdef WIN32
	typedef int socklen_t;
#else
    #include <sys/ioctl.h>
    #include <ctype.h>
	#include <errno.h>
	#include <unistd.h>
	#include <netdb.h>
	#define closesocket(x) close(x)
	#define ioctlsocket(a,b,c) ioctl(a,b,c)
	#define WSAGetLastError() (errno)
    #define WSAEWOULDBLOCK EINPROGRESS
#endif


//--------------------------------------------------------------------------------
bool TCPManager::m_bWinsockInit= false;

//--------------------------------------------------------------------------------
TCPManager::TCPManager()
{
	// was winsock initialized?
#ifdef WIN32
	if (!m_bWinsockInit) {
		WORD vr;
		WSADATA wsaData;
		vr= MAKEWORD(2, 2);
		WSAStartup(vr, &wsaData);
		m_bWinsockInit= true;
	}
#else
	m_bWinsockInit=	true;
#endif

	m_hSocket= INVALID_SOCKET;
	m_lTimeoutConnect= DEFAULT_TIMEOUT;
	m_lTimeoutAccept= DEFAULT_TIMEOUT;
	m_lTimeoutSend= DEFAULT_TIMEOUT;
	m_lTimeoutReceive= DEFAULT_TIMEOUT;
	m_iListenPort= -1;
	m_iMaxConnections = 0;
};

//--------------------------------------------------------------------------------
TCPManager::~TCPManager()
{
	if (m_hSocket != INVALID_SOCKET)
	{
		Close();
	}
};

//--------------------------------------------------------------------------------
/// Closes an open socket.
/// NOTE: A closed socket cannot be reused again without a call to "Create()".
bool TCPManager::Close()
{
	if (m_hSocket == INVALID_SOCKET)
	{
		return(false);
	}
	if(closesocket(m_hSocket) == SOCKET_ERROR)
	{
		return(false);
	}
	m_hSocket= INVALID_SOCKET;

	return(true);
}

void TCPManager::CleanUp()
{
	if (m_hSocket != INVALID_SOCKET)
	{
		Close();
	}
#ifdef WIN32
	WSACleanup();
#endif
	m_bWinsockInit= false;
}

//--------------------------------------------------------------------------------
bool TCPManager::Create()
{
	if (m_hSocket != INVALID_SOCKET)
	{
		return(false);
	}
	m_hSocket = socket(AF_INET, SOCK_STREAM, IPPROTO_IP);
	return(m_hSocket != INVALID_SOCKET);
}


//--------------------------------------------------------------------------------
bool TCPManager::Bind(USHORT usPort)
{
	if (m_hSocket == INVALID_SOCKET)
	{
		return(false);
	}

	struct sockaddr_in local;
	local.sin_family = AF_INET;
	local.sin_addr.s_addr = INADDR_ANY;
	//Port MUST be in Network Byte Order
	local.sin_port = htons(usPort);

	if (bind(m_hSocket,(struct sockaddr*)&local,sizeof(local)) == SOCKET_ERROR)
	{
		//LOG_ERROR((T("TCPManager::Bind: can't bind to port %u! Error: %d"), usPort, WSAGetLastError()));
		return false;
	}
	return true;
}

bool TCPManager::Listen(int iMaxConnections)
{
	if (m_hSocket == INVALID_SOCKET)
	{
		return(false);
	}
	m_iMaxConnections = iMaxConnections;
	return(listen(m_hSocket, iMaxConnections) != SOCKET_ERROR);
}

//--------------------------------------------------------------------------------
bool TCPManager::Accept(TCPManager& sConnect)
{
	if (m_hSocket == INVALID_SOCKET)
	{
		return(false);
	}

	if (m_lTimeoutAccept != NO_TIMEOUT)
	{
		fd_set fd;
		FD_ZERO(&fd);
		FD_SET(m_hSocket, &fd);
		struct timeval tv= { m_lTimeoutAccept, 0 };
		if(select((int)m_hSocket+1, &fd, NULL, NULL, &tv) <= 0)
		{
			return(false);
		}
	}

	sockaddr_in addr;
	socklen_t iSize= sizeof(addr);
	sConnect.m_hSocket= accept(m_hSocket, (sockaddr*)&addr, &iSize);
	return(sConnect.m_hSocket != INVALID_SOCKET);
}

//--------------------------------------------------------------------------------
//starts and stops blocking for the connect method
bool TCPManager::SetBlocking(bool bIsBlockingA)
{
	if (m_hSocket == INVALID_SOCKET)
	{
		return false;
	}

	unsigned long onL = (bIsBlockingA) ? 1:0;
	if (ioctlsocket(m_hSocket, FIONBIO, &onL) == SOCKET_ERROR)
	{
		//LOG_ERROR((T("TCPManager::SetBlocking Error: %d"), WSAGetLastError()));
		return false;
	}
	return true;
}

//--------------------------------------------------------------------------------
bool TCPManager::Connect(const char *pHost, USHORT usPort)
{
	if (m_hSocket == INVALID_SOCKET)
	{
		return(false);
	}

	sockaddr_in addr_in= {0};
	addr_in.sin_family= AF_INET; // host byte order
	addr_in.sin_port  = htons(usPort); // short, network byte order

	if (isalpha(pHost[0])) {  // host address is a name
		struct hostent *he;
		if ((he = gethostbyname(pHost)) == NULL)
		{
			//LOG_ERROR((T("TCPManager::Connect: gethostbyname(\"%hs\") failed!"), pHost));
			return false;
		}

		addr_in.sin_addr  = *((struct in_addr *)he->h_addr);
	}
	else {
		if ((addr_in.sin_addr.s_addr = inet_addr(pHost)) == INADDR_NONE)
		{
			//LOG_ERROR((T("TCPManager::Connect: inet_addr(\"%hs\") failed!"), pHost));
			return false;
		}
	}

	if (m_lTimeoutConnect != NO_TIMEOUT)
	{
		//set the blocking connect function to none blocking
		if (!SetBlocking(false))
		{
			return false;
		}
		if (connect(m_hSocket, (const struct sockaddr*)&addr_in, sizeof(addr_in)) == SOCKET_ERROR)
		{
			if (WSAGetLastError() != WSAEWOULDBLOCK)
			{
				//LOG_ERROR((T("TCPManager::Connect Error in nonblocking connect: %d"), WSAGetLastError()));
				return false;
			}
			// wait for connect using a select
			fd_set fd;
			FD_ZERO(&fd);
			FD_SET(m_hSocket, &fd);
			struct timeval tv= { m_lTimeoutConnect, 0 };
			int ret = select((int)m_hSocket+1, &fd, &fd, NULL, &tv);
			if (ret == 0)
			{
				//LOG_ERROR((T("TCPManager::Connect Timed out")));
				return false;
			}
			else if (ret < 0)
			{
				//LOG_ERROR((T("TCPManager::Connect Error in select: %d"), WSAGetLastError()));
				return false;
			}
		}

		// blocking socket
		SetBlocking(true);
	}
	else
	{
		if (connect(m_hSocket, (const struct sockaddr*)&addr_in, sizeof(addr_in)) == SOCKET_ERROR)
		{
			//LOG_ERROR((T("TCPManager::Connect Error: %d"), WSAGetLastError()));
			return false;
		}
	}

	return true;
}

//--------------------------------------------------------------------------------
int TCPManager::Write(const void* pBuff, const int iSize)
{
	int iBytesSent= 0;
	const char* pTemp= (const char*)pBuff;

	do {
		int iBytesTemp= Send(pTemp, iSize - iBytesSent);
		// error occured?
		if (iBytesTemp == SOCKET_ERROR)
		{
			return(SOCKET_ERROR);
		}
		if (iBytesTemp == SOCKET_TIMEOUT)
		{
			return(SOCKET_TIMEOUT);
		}

		iBytesSent+= iBytesTemp;
		pTemp+= iBytesTemp;
	} while(iBytesSent < iSize);

	return(iBytesSent);
}


//--------------------------------------------------------------------------------
/// Return values:
/// SOCKET_TIMEOUT indicates timeout
/// SOCKET_ERROR in case of a problem.
int TCPManager::Send(const void* pBuff, const int iSize)
{
	if (m_hSocket == INVALID_SOCKET)
	{
		return(SOCKET_ERROR);
	}

	if (m_lTimeoutSend != NO_TIMEOUT)
	{
		fd_set fd;
		FD_ZERO(&fd);
		FD_SET(m_hSocket, &fd);
		struct timeval tv= { m_lTimeoutSend, 0 };
		int ret = select((int)m_hSocket+1, NULL, &fd, NULL, &tv);
		if (ret == 0)
		{
			return(SOCKET_TIMEOUT);
		}
		else if (ret < 0)
		{
			return(SOCKET_ERROR);
		}
	}
	const char* pTemp= (const char*)pBuff;
	return(send(m_hSocket, pTemp, iSize, 0));
}

//--------------------------------------------------------------------------------
/// Return values:
/// SOCKET_TIMEOUT indicates timeout
/// SOCKET_ERROR in case of a problem.
int TCPManager::SendAll(const void* pBuff, const int iSize)
{
	if (m_hSocket == INVALID_SOCKET)
	{
		return(SOCKET_ERROR);
	}

	if (m_lTimeoutSend != NO_TIMEOUT)
	{
		fd_set fd;
		FD_ZERO(&fd);
		FD_SET(m_hSocket, &fd);
		struct timeval tv= { m_lTimeoutSend, 0 };
		int ret = select((int)m_hSocket+1, NULL, &fd, NULL, &tv);
		if (ret == 0)
		{
			return(SOCKET_TIMEOUT);
		}
		else if (ret < 0)
		{
			return(SOCKET_ERROR);
		}
	}

	const char* pTemp= (const char*)pBuff;
	int total= 0;
	int bytesleft = iSize;
	int ret= 0;

	while (total < iSize)
	{
		ret = send(m_hSocket, pTemp + total, bytesleft, 0);
		if (ret <= 0)
		{
			break;
		}
		total += ret;
		bytesleft -=ret;
	}

	return ret==-1?SOCKET_ERROR:total;
}


//--------------------------------------------------------------------------------
/// Return values:
/// SOCKET_TIMEOUT indicates timeout
/// SOCKET_ERROR in case of a problem.
int TCPManager::Receive(void* pBuff, const int iSize)
{
	if (m_hSocket == INVALID_SOCKET)
	{
		return(SOCKET_ERROR);
	}

	if (m_lTimeoutReceive != NO_TIMEOUT)
	{
		fd_set fd;
		FD_ZERO(&fd);
		FD_SET(m_hSocket, &fd);
		struct timeval tv= { m_lTimeoutReceive, 0 };
		int ret = select((int)m_hSocket+1, &fd, NULL, NULL, &tv);
		if (ret == 0)
		{
			return(SOCKET_TIMEOUT);
		}
		else if (ret < 0)
		{
			return(SOCKET_ERROR);
		}
	}
	char* pTemp= (char*)pBuff;
	return(recv(m_hSocket, pTemp, iSize, 0));
}


//--------------------------------------------------------------------------------
/// Return values:
/// SOCKET_TIMEOUT indicates timeout
/// SOCKET_ERROR in case of a problem.
int TCPManager::ReceiveAll(void* pBuff, const int iSize)
{
	if (m_hSocket == INVALID_SOCKET)
	{
		return(SOCKET_ERROR);
	}

	if (m_lTimeoutReceive != NO_TIMEOUT)
	{
		fd_set fd;
		FD_ZERO(&fd);
		FD_SET(m_hSocket, &fd);
		struct timeval tv= { m_lTimeoutReceive, 0 };
		int ret = select((int)m_hSocket+1, &fd, NULL, NULL, &tv);
		if (ret == 0)
		{
			return(SOCKET_TIMEOUT);
		}
		else if (ret < 0)
		{
			return(SOCKET_ERROR);
		}
	}

	char* pTemp= (char*)pBuff;
	int totalBytes=0;

	do {
		int ret= recv(m_hSocket, pTemp+totalBytes, iSize-totalBytes, 0);
		if (ret==0 && totalBytes != iSize)
		{
			return SOCKET_ERROR;
		}
		if (ret < 0)
		{
			return SOCKET_ERROR;
		}
		totalBytes += ret;
	} while(totalBytes < iSize);

	return totalBytes;
}

//--------------------------------------------------------------------------------
bool TCPManager::GetRemoteAddr(LPINETADDR pInetAddr)
{
	if (m_hSocket == INVALID_SOCKET)
	{
		return(false);
	}

	socklen_t iSize= sizeof(*pInetAddr);
	return(getpeername(m_hSocket, *pInetAddr, &iSize) != SOCKET_ERROR);
}

//--------------------------------------------------------------------------------
bool TCPManager::GetInetAddr(LPINETADDR pInetAddr)
{
	if (m_hSocket == INVALID_SOCKET)
	{
		return(false);
	}

	socklen_t iSize= sizeof(*pInetAddr);
	return(getsockname(m_hSocket, *pInetAddr, &iSize) != SOCKET_ERROR);
}

int TCPManager::GetReceiveBufferSize()
{
	if (m_hSocket == INVALID_SOCKET)
	{
		return(false);
	}

	int sizeBuffer=0;
	socklen_t size = sizeof(sizeBuffer);
	getsockopt(m_hSocket, SOL_SOCKET, SO_RCVBUF, (char*)&sizeBuffer, &size);

	return sizeBuffer;
}

bool TCPManager::SetReceiveBufferSize(int sizeInByte)
{
	if (m_hSocket == INVALID_SOCKET)
	{
		return(false);
	}

	return (setsockopt(m_hSocket, SOL_SOCKET, SO_RCVBUF, (char*)&sizeInByte, sizeof(sizeInByte)) != SOCKET_ERROR);
}

int TCPManager::GetSendBufferSize()
{
	if (m_hSocket == INVALID_SOCKET)
	{
		return(false);
	}

	int sizeBuffer=0;
	socklen_t size = sizeof(sizeBuffer);
	getsockopt(m_hSocket, SOL_SOCKET, SO_SNDBUF, (char*)&sizeBuffer, &size);

	return sizeBuffer;
}

bool TCPManager::SetSendBufferSize(int sizeInByte)
{
	if (m_hSocket == INVALID_SOCKET)
	{
		return(false);
	}

	return (setsockopt(m_hSocket, SOL_SOCKET, SO_SNDBUF, (char*)&sizeInByte, sizeof(sizeInByte)) != SOCKET_ERROR);
}

bool TCPManager::CheckHost(const char *pAddrStr)
{
	if (isalpha(pAddrStr[0]))  // host address is a name
	{
		return (gethostbyname(pAddrStr) != NULL);
	}
	else
	{
		return (inet_addr(pAddrStr) != INADDR_NONE);
	}
}
